#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
@Project : ailearning
@File : rag_base_retrieve.py
@Author : 李峰
@Date : 2025/9/23 17:41
@Software: PyCharm
说明：RAG 基础步骤 检索（以web网站为例）
加载网站》分割文档》嵌入向量库》检索
'''
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader

llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,
                 api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")
# --------------简单的一段话-----------------------
# 1、加载文本
loader = TextLoader(file_path="xiyouji.txt",encoding="utf-8")
blog_docs = loader.load()
# 2、分割文档
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=300,chunk_overlap=50)
splits = text_splitter.split_documents(blog_docs)
# 3、构建一个支持语义搜索的向量数据库（数据库在内存中）
embd=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
vectorstore = Chroma.from_documents(documents=splits,embedding=embd)
#4、检索
# 将一个向量数据库（vectorstore）转换为一个检索器（retriever），并设置检索参数：返回最相似的5个结果。
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
docs = retriever.get_relevant_documents("孙悟空的出生地在哪里？")
print(f"打印检索到的结果：{docs[0]}\n{docs[1]}\n{docs[2]}\n{docs[3]}\n{docs[4]}")